/* `include "define.sv"
// 乘法器使能信号始终保持高电平
module pe (
    input clk,
    input reset,                  // 高电平复位
    input [1:0] data_type,        // 数据类型（INT4/INT8/FP16/FP32等）
    input mult_en,                // 乘法器使能
    input c_add_en,                // C矩阵加法使能信号
    input [`DATA_WIDTH-1:0] a_in, // 输入A（来自左侧PE）
    input [`DATA_WIDTH-1:0] b_in, // 输入B（来自上侧PE）
    input [`DATA_WIDTH-1:0] c_in,  // 输入C的数据
    output reg [`DATA_WIDTH-1:0] a_out, // 向右传递A（延迟一拍）
    output reg [`DATA_WIDTH-1:0] b_out, // 向下传递B（延迟一拍）
    output reg [`DATA_WIDTH-1:0] c_out,  // 累加结果输出
    input mix_precision,
    output add_valid_o,          // 流向下一级的valid
    output reg calc_done,         // 计算结束信号
    output reg mult_en_out,
    input pe_reset                // PE内部复位信号，用于重置控制逻辑
);

// 内部信号声明
wire [`DATA_WIDTH-1:0] mult_out; // 乘法器输出结果
wire mult_valid;                 // 乘法器输出有效信号
wire [`DATA_WIDTH-1:0] add_out; // 加法器输出结果
wire adder_val_o;

reg c_added; // C矩阵添加标志
reg calc_started; // 计算开始标志
reg [1:0] data_type_mix;
reg [`DATA_WIDTH-1:0] mult_out_reg; //乘法器输出寄存一拍
reg mult_valid_r; //输出信号寄存

// 累加器寄存器（保存历史累加结果）
reg [`DATA_WIDTH-1:0] accumulator;  //历史累加结果
//reg [`DATA_WIDTH-1:0] op2_in;       //来自乘法器结果或C矩阵寄存器的值
//reg [`DATA_WIDTH-1:0] c_reg;   // C矩阵数值寄存器

// 数据传递逻辑
always @(posedge clk or posedge reset) begin
    if (reset) begin
        a_out <= {`DATA_WIDTH{1'b0}};
        b_out <= {`DATA_WIDTH{1'b0}};
        mult_en_out <= 1'b0;
        mult_out_reg <= 32'b0;
        mult_valid_r <= 1'b0;
    end else begin
        a_out <= a_in;  
        b_out <= b_in;  
        mult_en_out <= mult_en;
        mult_out_reg <= mult_out;
        mult_valid_r <= mult_valid;
    end
end

always @(*) begin
    if(mix_precision) begin
        data_type_mix = `FP32;
    end
    else begin
        data_type_mix = data_type;
    end
end
//  乘法器模块
Mult_top  mult_inst (
    .clk(clk),
    .sys_rst_n(!reset),
    .en(mult_en),               // 乘法器使能
    .data_type(data_type),
    .mix_precision(mix_precision),
    .data_in(a_in),          // 直接使用当前输入（非寄存后的值）
    .weight_in(b_in),        // 直接使用当前输入（非寄存后的值）
    .float_o(mult_out),
    .mult_out_valid(mult_valid)
);

//  加法器模块（累加乘法结果）
Adder_top add_inst (
    .clk(clk),
    .sys_rst_n(!reset),
    .valid1(calc_started),          // 属于累加的valid信号
    .valid2(mult_valid||mult_valid_r),    // 乘法结果有效时触发加法
    .data_type(data_type_mix),
    .a_in(accumulator),     // 历史累加结果（操作数1）
    .b_in(mult_out_reg),        // 乘法结果（操作数2）
    .mix_precision(mix_precision),
    .float_o(add_out),
    .adder_valid(adder_val_o)
);

// 累加器更新逻辑
reg add_valid_prev; // 存储add_valid_o的前一状态
assign add_valid_o = adder_val_o&&mult_valid_r;
always @(posedge clk or posedge reset) begin
    if (reset) begin
        c_added      <= 1'b1;
        accumulator  <= {`DATA_WIDTH{1'b0}};
        c_out        <= {`DATA_WIDTH{1'b0}};
        calc_done    <= 1'b0;
        calc_started <= 1'b0;
        add_valid_prev <= 1'b0; // 复位时初始化
    end else begin
        add_valid_prev <= add_valid_o; // 保存前一状态

        // PE内部复位逻辑：重置控制信号但保持计算结果
        if (pe_reset) begin
            c_added      <= 1'b1;        // 重置C矩阵添加标志
            calc_done    <= 1'b0;        // 清除计算完成标志
            calc_started <= 1'b0;        // 清除计算开始标志
            add_valid_prev <= 1'b0;      // 重置valid信号状态
            // 注意：不重置accumulator和c_out，保持最终计算结果
        end
        // C矩阵加载逻辑
        else if (c_add_en & c_added) begin
            c_added     <= 1'b0;
            accumulator <= c_in;
            calc_done   <= 1'b0;
            calc_started <= 1'b1;
        end 
        // 正常累加逻辑
        else if (add_valid_o) begin
            accumulator <= add_out;
            c_out       <= add_out;
        end

        // 检测add_valid_o的下降沿（从1变0）
        if (add_valid_prev && !add_valid_o) begin
            if (calc_started) begin
                calc_done <= 1'b1;       // 设置计算结束标志
                calc_started <= 1'b0;    // 清除计算开始标志
            end
        end
    end
end

endmodule

 */


 `include "define.sv"

module pe (
    input clk,
    input reset,                  // 高电平复位
    input [1:0] data_type,        // 数据类型（INT4/INT8/FP16/FP32等）
    input mult_en,                // 乘法器使能
    input c_add_en,               // C矩阵加法使能信号
    input [`DATA_WIDTH-1:0] a_in, // 输入A（来自左侧PE）
    input [`DATA_WIDTH-1:0] b_in, // 输入B（来自上侧PE）
    input [`DATA_WIDTH-1:0] c_in, // 输入C的数据
    output reg [`DATA_WIDTH-1:0] a_out, // 向右传递A（延迟一拍）
    output reg [`DATA_WIDTH-1:0] b_out, // 向下传递B（延迟一拍）
    output reg [`DATA_WIDTH-1:0] c_out, // 累加结果输出
    input mix_precision,
    output reg add_valid_o,       // 流向下一级的valid
    output reg calc_done,         // 计算结束信号
    output reg mult_en_out,
    input pe_reset                // PE内部复位信号，用于重置控制逻辑
);

// 状态机定义
typedef enum logic [1:0] {
    IDLE = 2'b00,      // 等待数据输入
    CALCULATE = 2'b01, // 数据计算中
    DONE = 2'b10       // 计算结束，数据正在被下一级读取
} pe_state_t;

pe_state_t current_state, next_state;

// 内部信号声明
wire [`DATA_WIDTH-1:0] mult_out; // 乘法器输出结果
wire mult_valid;                  // 乘法器输出有效信号
wire [`DATA_WIDTH-1:0] add_out;  // 加法器输出结果
wire adder_val_o;

reg c_added;                      // C矩阵添加标志
reg [1:0] data_type_mix;
reg [`DATA_WIDTH-1:0] mult_out_reg; // 乘法器输出寄存一拍
reg mult_valid_r;                 // 输出信号寄存
reg [`DATA_WIDTH-1:0] accumulator;  // 历史累加结果
reg add_valid_prev;               // 存储add_valid_o的前一状态

// 状态机时序逻辑
always @(posedge clk or posedge reset) begin
    if (reset) begin
        current_state <= IDLE;
    end else begin
        current_state <= next_state;
    end
end

// 状态机组合逻辑
always @(*) begin
    case (current_state)
        IDLE: begin
            if (c_add_en) begin
                next_state = CALCULATE;
            end else begin
                next_state = IDLE;
            end
        end
        
        CALCULATE: begin
            // 检测add_valid_o的下降沿（从1变0）
            if (add_valid_prev && !add_valid_o) begin
                next_state = DONE;
            end else begin
                next_state = CALCULATE;
            end
        end
        
        DONE: begin
            if (pe_reset) begin
                next_state = IDLE;
            end else begin
                next_state = DONE;
            end
        end
        
        default: begin
            next_state = IDLE;
        end
    endcase
end

// 数据传递逻辑
always @(posedge clk or posedge reset) begin
    if (reset) begin
        a_out <= {`DATA_WIDTH{1'b0}};
        b_out <= {`DATA_WIDTH{1'b0}};
        mult_en_out <= 1'b0;
        mult_out_reg <= 32'b0;
        mult_valid_r <= 1'b0;
    end else begin
        a_out <= a_in;  
        b_out <= b_in;  
        mult_en_out <= mult_en;
        mult_out_reg <= mult_out;
        mult_valid_r <= mult_valid;
    end
end

// 数据类型选择逻辑
always @(*) begin
    if(mix_precision) begin
        data_type_mix = `FP32;
    end else begin
        data_type_mix = data_type;
    end
end

// 乘法器模块
Mult_top mult_inst (
    .clk(clk),
    .sys_rst_n(!reset),
    .en(mult_en),               // 乘法器使能
    .data_type(data_type),
    .mix_precision(mix_precision),
    .data_in(a_in),          // 直接使用当前输入（非寄存后的值）
    .weight_in(b_in),        // 直接使用当前输入（非寄存后的值）
    .float_o(mult_out),
    .mult_out_valid(mult_valid)
);

// 加法器模块（累加乘法结果）
Adder_top add_inst (
    .clk(clk),
    .sys_rst_n(!reset),
    .valid1(current_state == CALCULATE),  // 在CALCULATE状态下才进行累加
    .valid2(mult_valid || mult_valid_r),  // 乘法结果有效时触发加法
    .data_type(data_type_mix),
    .a_in(accumulator),     // 历史累加结果（操作数1）
    .b_in(mult_out_reg),    // 乘法结果（操作数2）
    .mix_precision(mix_precision),
    .float_o(add_out),
    .adder_valid(adder_val_o)
);

// 累加器更新逻辑
assign add_valid_o = adder_val_o && mult_valid_r && (current_state == CALCULATE);

always @(posedge clk or posedge reset) begin
    if (reset) begin
        c_added <= 1'b1;
        accumulator <= {`DATA_WIDTH{1'b0}};
        c_out <= {`DATA_WIDTH{1'b0}};
        calc_done <= 1'b0;
        add_valid_prev <= 1'b0;
    end else begin
        add_valid_prev <= add_valid_o; // 保存前一状态

        case (current_state)
            IDLE: begin
                // 在IDLE状态下，等待C矩阵加载
                if (c_add_en && c_added) begin
                    c_added <= 1'b0;
                    accumulator <= c_in;
                    calc_done <= 1'b0;
                end
            end
            
            CALCULATE: begin
                // 在CALCULATE状态下，进行累加计算
                if (add_valid_o) begin
                    accumulator <= add_out;
                    c_out <= add_out;
                end
            end
            
            DONE: begin
                // 在DONE状态下，设置计算完成标志
                calc_done <= 1'b1;
            end
        endcase

        // PE复位处理：重置控制信号但保持计算结果
        if (pe_reset) begin
            c_added <= 1'b1;        // 重置C矩阵添加标志
            calc_done <= 1'b0;      // 清除计算完成标志
            add_valid_prev <= 1'b0; // 重置valid信号状态
            // 注意：不重置accumulator和c_out，保持最终计算结果
        end
    end
end

/* // 调试输出（可选）
// synthesis translate_off
always @(posedge clk) begin
    if (current_state != next_state) begin
        $display("[%t] PE State Change: %s -> %s", $time, 
                 current_state.name(), next_state.name());
    end
end
// synthesis translate_on */

endmodule 